/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <fmt/format.h>

#include <mariana-trench/JsonValidation.h>
#include <mariana-trench/KindFactory.h>
#include <mariana-trench/NamedKind.h>
#include <mariana-trench/Rule.h>
#include <mariana-trench/RulesCoverage.h>
#include <mariana-trench/SourceSinkWithExploitabilityRule.h>
#include <mariana-trench/TransformsFactory.h>

namespace marianatrench {

bool SourceSinkWithExploitabilityRule::uses(const Kind* kind) const {
  const auto base_kind = kind->discard_transforms();
  return effect_source_kinds_.count(base_kind) != 0 ||
      source_kinds_.count(base_kind) != 0 || sink_kinds_.count(base_kind) != 0;
}

const TransformList* MT_NULLABLE
SourceSinkWithExploitabilityRule::source_as_transform(
    const Kind* source_kind) const {
  auto transform = source_as_transforms_.find(source_kind);

  if (transform == source_as_transforms_.end()) {
    return nullptr;
  }

  return transform->second;
}

std::optional<CoveredRule> SourceSinkWithExploitabilityRule::coverage(
    const KindSet& sources,
    const KindSet& sinks,
    const TransformSet& /* transforms */) const {
  auto used_rule_effect_sources =
      Rule::intersecting_kinds(/* rule_kinds */ effect_source_kinds(), sources);
  if (used_rule_effect_sources.empty()) {
    return std::nullopt;
  }

  auto used_rule_sources =
      Rule::intersecting_kinds(/* rule_kinds */ source_kinds(), sources);
  if (used_rule_sources.empty()) {
    return std::nullopt;
  }

  auto used_rule_sinks =
      Rule::intersecting_kinds(/* rule_kinds */ sink_kinds(), sinks);
  if (used_rule_sinks.empty()) {
    return std::nullopt;
  }

  used_rule_sources.merge(used_rule_effect_sources);

  return CoveredRule{
      .code = code(),
      .used_sources = std::move(used_rule_sources),
      .used_sinks = std::move(used_rule_sinks),
      .used_transforms = TransformSet{},
  };
}

std::unique_ptr<Rule> SourceSinkWithExploitabilityRule::from_json(
    const std::string& name,
    int code,
    const std::string& description,
    const Json::Value& value,
    Context& context) {
  JsonValidation::check_unexpected_members(
      value,
      {"name",
       "code",
       "description",
       "effect_sources",
       "sources",
       "sinks",
       "oncall"});

  KindSet effect_source_kinds;
  for (const auto& effect_source_kind :
       JsonValidation::nonempty_array(value, /* field */ "effect_sources")) {
    effect_source_kinds.insert(
        NamedKind::from_json(effect_source_kind, context));
  }

  KindSet source_kinds;
  KindToTransformsMap source_as_transforms;
  for (const auto& source_kind_value :
       JsonValidation::nonempty_array(value, /* field */ "sources")) {
    const auto* source_kind = NamedKind::from_json(source_kind_value, context);
    source_kinds.insert(source_kind);
    source_as_transforms.emplace(
        source_kind,
        context.transforms_factory->create(
            TransformList::from_kind(source_kind, context)));
  }

  KindSet sink_kinds;
  for (const auto& sink_kind :
       JsonValidation::nonempty_array(value, /* field */ "sinks")) {
    sink_kinds.insert(NamedKind::from_json(sink_kind, context));
  }

  return std::make_unique<SourceSinkWithExploitabilityRule>(
      name,
      code,
      description,
      effect_source_kinds,
      std::move(source_kinds),
      std::move(sink_kinds),
      std::move(source_as_transforms));
}

Json::Value SourceSinkWithExploitabilityRule::to_json() const {
  auto value = Rule::to_json();

  auto effect_source_kinds_value = Json::Value(Json::arrayValue);
  for (const auto* effect_source_kind : effect_source_kinds_) {
    effect_source_kinds_value.append(effect_source_kind->to_json());
  }

  auto source_kinds_value = Json::Value(Json::arrayValue);
  for (const auto* source_kind : source_kinds_) {
    source_kinds_value.append(source_kind->to_json());
  }

  auto sink_kinds_value = Json::Value(Json::arrayValue);
  for (const auto* sink_kind : sink_kinds_) {
    sink_kinds_value.append(sink_kind->to_json());
  }

  value["effect_sources"] = effect_source_kinds_value;
  value["sources"] = source_kinds_value;
  value["sinks"] = sink_kinds_value;

  return value;
}

} // namespace marianatrench
